import cv2
import numpy as np
import serial
import time

# —-------------------- 初始化串口 --------------------
ser = serial.Serial('/dev/ttyUSB0', baudrate=115200, timeout=0.5) # 串口号自己查询
time.sleep(2)

# -------------------- HSV 阈值设置（蓝紫激光示例） --------------------
HSV_LOWER = np.array([0, 0, 214])
HSV_UPPER = np.array([41, 28, 255])

# -------------------- 是否为矩形判断 --------------------
def is_rectangle(approx, angle_threshold=40):
    if len(approx) != 4:
        return False
    def angle(pt1, pt2, pt0):
        dx1, dy1 = pt1[0][0] - pt0[0][0], pt1[0][1] - pt0[0][1]
        dx2, dy2 = pt2[0][0] - pt0[0][0], pt2[0][1] - pt0[0][1]
        return (dx1 * dx2 + dy1 * dy2) / np.sqrt((dx1**2 + dy1**2) * (dx2**2 + dy2**2))
    for i in range(4):
        cos = angle(approx[i], approx[(i + 2) % 4], approx[(i + 1) % 4])
        deg = np.arccos(cos) * 180 / np.pi
        if abs(deg - 90) > angle_threshold:
            return False
    return True

# -------------------- 激光点检测 --------------------
def detect_laser_points(frame):
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(hsv, HSV_LOWER, HSV_UPPER)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8))
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    for cnt in contours:
        area = cv2.contourArea(cnt)
        if 2 < area < 200:
            (x, y), _ = cv2.minEnclosingCircle(cnt)
            return int(x), int(y)
    return None

# -------------------- ROI 区域中矩形中心提取 --------------------
def process_roi_area(frame, roi):
    x1, y1, x2, y2 = roi
    roi_img = frame[y1:y2, x1:x2]

    gray = cv2.cvtColor(roi_img, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (9, 9), 3)
    edges = cv2.Canny(blurred, 50, 150)
    edges = cv2.dilate(edges, None, iterations=2)
    edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8))

    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    rect_center = None

    for cnt in contours:
        approx = cv2.approxPolyDP(cnt, 0.02 * cv2.arcLength(cnt, True), True)
        area = cv2.contourArea(cnt)
        if area < 3000 or area > 50000:
            continue
        if not cv2.isContourConvex(approx) or not is_rectangle(approx):
            continue

        approx += np.array([[[x1, y1]]])  # 偏移回原图坐标
        corners = approx.reshape(4, 2)
        rect_center = tuple(np.mean(corners, axis=0).astype(int))

        # 画图
        for pt in corners:
            cv2.circle(frame, tuple(pt), 4, (0, 0, 255), -1)
        cv2.circle(frame, rect_center, 4, (0, 0, 255), -1)
        cv2.putText(frame, f"Center: {rect_center}", (rect_center[0]+10, rect_center[1]),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
        cv2.drawContours(frame, [approx], -1, (0, 255, 0), 2)

    cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 1)
    return frame, edges, rect_center

# -------------------- 串口数据打包发送（帧头 0xAA55，帧尾 0x0D0A）--------------------
def send_serial_packet(laser, rect):
    if not laser or not rect:
        return

    def clamp_and_pack(x):
        return max(0, min(x, 65535)).to_bytes(2, 'big')

    packet = b'\xAA\x55'  # 帧头
    packet += clamp_and_pack(laser[0]) + clamp_and_pack(laser[1])
    packet += clamp_and_pack(rect[0]) + clamp_and_pack(rect[1])
    packet += b'\x0D\x0A'  # 帧尾

    try:
        ser.write(packet)
    except Exception as e:
        print("[串口发送失败]", e)


# -------------------- 主函数：视频流检测 --------------------
def run_video_detection(source=1, roi_coords=(50, 50, 590, 380)):  # 默认摄像头1，ROI区域坐标
    cap = cv2.VideoCapture(source)
    if not cap.isOpened():
        print("[错误] 无法打开视频流")
        return

    # 初始化卡尔曼滤波器（这个目前写了，暂时没有用只是有效果，发送的数据还是真实矩形中心坐标，后续如果小车晃动非常厉害识别效果不好再用预测坐标）
    kalman = cv2.KalmanFilter(4, 2)
    kalman.measurementMatrix = np.array([[1, 0, 0, 0],
                                         [0, 1, 0, 0]], dtype=np.float32)
    kalman.transitionMatrix = np.array([[1, 0, 1, 0],
                                        [0, 1, 0, 1],
                                        [0, 0, 1, 0],
                                        [0, 0, 0, 1]], dtype=np.float32)
    kalman.processNoiseCov = np.eye(4, dtype=np.float32) * 0.03

    miss_count = 0
    max_miss_count = 10

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        laser_point = detect_laser_points(frame)
        frame, edges, rect_center = process_roi_area(frame, roi_coords)

        # 卡尔曼滤波更新
        if rect_center is not None:
            kalman.correct(np.array([[np.float32(rect_center[0])],
                                     [np.float32(rect_center[1])]]))
            miss_count = 0
        else:
            miss_count += 1
            if miss_count > max_miss_count:
                kalman.statePre[:] = 0
                kalman.statePost[:] = 0
                miss_count = 0
                print("[INFO] 卡尔曼状态重置")

        predicted = kalman.predict()
        predicted_point = (int(predicted[0]), int(predicted[1])) # 这里是预测的卡尔曼坐标
        cv2.circle(frame, predicted_point, 3, (0, 255, 255), -1)
        cv2.putText(frame, f"Kalman: {predicted_point}", (predicted_point[0]+10, predicted_point[1]),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 200, 200), 1)

        if laser_point:
            cv2.circle(frame, laser_point, 2, (255, 0, 0), -1)
            cv2.putText(frame, f"Laser: {laser_point}", (laser_point[0]+10, laser_point[1]),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 255), 1)

        if laser_point and rect_center:
            send_serial_packet(laser_point, rect_center) # 打包并发送数据（实际矩形中心坐标和激光坐标，如果要修改为预测的卡尔曼坐标，把rect_center改成predicted_point即可）
            print(f"Laser: {laser_point}, Rect: {rect_center}")

        cv2.imshow("ROI Rectangle Detection", frame)
        cv2.imshow("Edge View", edges)
        if cv2.waitKey(1) & 0xFF == 27:  # ESC键退出
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    run_video_detection()
